# backend/main.py
import os
import time
import json
import base64
import logging
import asyncio
import numpy as np
from typing import Generator
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Body
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles

from db import init_db, redis_client
from asr import StreamingASR
from llm import llm_stream_generate
from tts import synth_tts_bytes, synth_tts_savefile
from utils import detect_keyword, save_llm_fragment

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("main")

app = FastAPI()
# mount audio dir to serve saved files
AUDIO_DIR = os.getenv("AUDIO_DIR", "./audio")
os.makedirs(AUDIO_DIR, exist_ok=True)
app.mount("/audio", StaticFiles(directory=AUDIO_DIR), name="audio")

# init db
init_db()

# ASR instance
asr = StreamingASR()

# fixed conv id
CONV_ID = 1

# mapping from session_id -> ws connection state if needed
WS_SESSIONS = {}

@app.websocket("/ws/asr")
async def ws_asr(ws: WebSocket):
    await ws.accept()
    sid = str(id(ws))
    session_key = f"session_{sid}"
    asr.start(session_key)
    WS_SESSIONS[session_key] = {"ws": ws}
    logger.info("WS connected %s", session_key)

    try:
        while True:
            # receive binary audio bytes from client (16-bit PCM)
            data = await ws.receive_bytes()
            # convert to int16 numpy
            pcm = np.frombuffer(data, dtype=np.int16)
            # get partial
            partial = asr.accept_chunk(session_key, pcm, is_final=False)
            if partial:
                # send partial to client
                await ws.send_json({"event": "asr_partial", "text": partial})
                k = detect_keyword(partial)
                if k == "interrupt":
                    # interrupt: notify client and reset ASR session
                    await ws.send_json({"event": "interrupt", "reason": "user_interrupt"})
                    asr.finalize(session_key)
                    asr.start(session_key)
                    continue
                elif k == "wake":
                    await ws.send_json({"event": "wake", "text": partial})
    except WebSocketDisconnect:
        logger.info("WS disconnect %s", session_key)
    except Exception:
        logger.exception("WS error")
    finally:
        try:
            await ws.close()
        except:
            pass
        asr.sessions.pop(session_key, None)
        WS_SESSIONS.pop(session_key, None)
        logger.info("WS closed %s", session_key)

# SSE streaming endpoint: client posts session_id, optionally conv_id and save_audio
@app.post("/finalize_stream")
async def finalize_stream_endpoint(request: Request, session_id: str = Body(...), conv_id: int = Body(None), save_audio: bool = Body(False)):
    """
    Finalize ASR for session_id, then stream LLM-generated fragments via SSE.
    conv_id: if provided and save_audio True, fragments are saved to DB and audio files kept.
    save_audio: boolean, if True backend saves per-fragment wav files and returns audio_url; otherwise returns audio_base64.
    """
    session_prefix = f"session_{session_id}"
    try:
        final_text = asr.finalize(session_prefix)
    except Exception:
        final_text = ""

    if not final_text:
        def _no_speech() -> Generator[bytes, None, None]:
            yield _sse_pack("no_speech", {"text": ""})
        return StreamingResponse(_no_speech(), media_type="text/event-stream")

    # require wake word
    if "小爱" not in final_text:
        def _ignored():
            yield _sse_pack("ignored", {"text": final_text})
        return StreamingResponse(_ignored(), media_type="text/event-stream")

    # normalize: remove leading wake words
    normalized = final_text
    while normalized.startswith("小爱"):
        normalized = normalized[len("小爱"):].strip()

    def event_generator():
        seg_idx = 0
        for chunk in llm_stream_generate(normalized):
            text_chunk = chunk.get("text", "") or ""
            if not text_chunk:
                continue

            # handle interrupt inside generation
            if detect_keyword(text_chunk) == "interrupt":
                yield _sse_pack("interrupt", {"reason": "llm_fragment_interrupt"})
                break

            if save_audio:
                # save file
                filepath = synth_tts_savefile(text_chunk)
                filename = os.path.basename(filepath) if filepath else ""
                audio_url = f"/audio/{filename}" if filename else ""
                payload = {"text": text_chunk, "audio_url": audio_url, "segment_index": seg_idx}
            else:
                audio_bytes = synth_tts_bytes(text_chunk)
                b64 = "" if audio_bytes is None else base64.b64encode(audio_bytes).decode("ascii")
                payload = {"text": text_chunk, "audio_base64": b64, "segment_index": seg_idx}

            # persist if conv_id and save_audio (you can adjust to persist even without audio)
            if conv_id is not None and save_audio:
                try:
                    save_llm_fragment(conv_id, "assistant", text_chunk, filepath, seg_idx)
                except Exception:
                    logger.exception("saving fragment failed")

            yield _sse_pack("fragment", payload)
            seg_idx += 1
            # slight throttle
            time.sleep(0.01)

        yield _sse_pack("done", {"segments": seg_idx})

    return StreamingResponse(event_generator(), media_type="text/event-stream")

def _sse_pack(event: str, data: dict) -> bytes:
    payload = {"event": event, "data": data}
    s = f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
    return s.encode("utf-8")

# optionally simple health or other endpoints
@app.get("/")
async def root():
    return {"message": "voice realtime demo running", "conv_id": CONV_ID}
